// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/values.h>
#include <pollux/common/testutil/test_value.h>

using kumo::pollux::common::testutil::TestValue;

namespace kumo::pollux::exec {
    Values::Values(
        int32_t operatorId,
        DriverCtx *driverCtx,
        std::shared_ptr<const core::ValuesNode> values)
        : SourceOperator(
              driverCtx,
              values->outputType(),
              operatorId,
              values->id(),
              "Values"),
          valueNodes_(std::move(values)),
          roundsLeft_(valueNodes_->repeatTimes()) {
    }

    void Values::initialize() {
        Operator::initialize();
        POLLUX_CHECK_NOT_NULL(valueNodes_);
        POLLUX_CHECK(values_.empty());
        // Drop empty vectors. Operator::getOutput is expected to return nullptr or a
        // non-empty vector.
        values_.reserve(valueNodes_->values().size());
        for (auto &vector: valueNodes_->values()) {
            if (vector->size() > 0) {
                if (valueNodes_->isParallelizable()) {
                    // If this is parallelizable, copy the values to prevent Vectors from
                    // being shared across threads.  Note that the contract in ValuesNode is
                    // that this should only be enabled for testing.
                    values_.emplace_back(std::static_pointer_cast<RowVector>(
                        vector->copy_preserve_encodings()));
                } else {
                    values_.emplace_back(vector);
                }
            }
        }
        // Drop the reference on the value nodes.
        valueNodes_ = nullptr;
    }

    RowVectorPtr Values::getOutput() {
        TestValue::adjust("kumo::pollux::exec::Values::getOutput", this);
        if (current_ >= values_.size()) {
            if (roundsLeft_ > 0) {
                --roundsLeft_;
            }
            current_ = 0;
        }

        if (isFinished()) {
            return nullptr;
        }
        return values_[current_++];
    }

    void Values::close() {
        current_ = values_.size();
        roundsLeft_ = 0;
    }

    bool Values::isFinished() {
        return roundsLeft_ == 0;
    }
} // namespace kumo::pollux::exec
